import os
import random

def create_calibration_dataset_txt(base_data_dir, subset_folder, num_images, output_file='dataset.txt'):
    """
    从指定的数据集文件夹中随机选择图片路径，并写入到 dataset.txt 文件中。

    Args:
        base_data_dir (str): 数据集根目录，例如 'iflytek_phone_ucar'。
                             脚本会从这个目录下的 'images/subset_folder' 中查找图片。
        subset_folder (str): 要选择图片的子文件夹，例如 'train' 或 'val'。
        num_images (int): 要随机选择的图片数量。
        output_file (str): 输出的 dataset.txt 文件名。
    """
    # 构建包含图片的完整路径，例如 'iflytek_phone_ucar/images/val'
    full_image_dir = os.path.join(base_data_dir, 'images', subset_folder)

    if not os.path.isdir(full_image_dir):
        print(f"错误：找不到图片目录 '{full_image_dir}'。请检查 'YOUR_DATASET_ROOT' 和 'SUBSET_TO_USE' 配置是否正确。")
        return

    all_image_paths = []
    # 定义常见图片扩展名
    image_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.webp', '.tiff')

    print(f"正在扫描图片目录：'{full_image_dir}'...")
    for root, _, files in os.walk(full_image_dir):
        for file in files:
            # 检查文件是否是图片类型
            if file.lower().endswith(image_extensions):
                # 构建相对于脚本当前工作目录的图片路径
                # 这要求 base_data_dir 也是相对于当前工作目录的
                relative_path = os.path.relpath(os.path.join(root, file), start=os.getcwd())
                all_image_paths.append(relative_path)

    if not all_image_paths:
        print(f"警告：在 '{full_image_dir}' 中没有找到任何图片文件。请确保文件夹中有图片。")
        return

    print(f"共找到 {len(all_image_paths)} 张图片。")

    if num_images > len(all_image_paths):
        print(f"警告：请求选择的图片数量 ({num_images}) 超过了可用图片数量 ({len(all_image_paths)})。将选择所有可用图片。")
        selected_images = all_image_paths
    else:
        # 随机选择指定数量的图片
        selected_images = random.sample(all_image_paths, num_images)

    # 将选中的图片路径写入到输出文件
    with open(output_file, 'w') as f:
        for img_path in selected_images:
            f.write(img_path + '\n')

    print(f"成功从 '{full_image_dir}' 中随机选择了 {len(selected_images)} 张图片，并写入到 '{output_file}'。")
    print(f"这些路径是相对于脚本运行目录的，请确保 '{output_file}' 文件在您运行 RKNN 转换脚本的同一目录下。")

if __name__ == '__main__':
    # --- 用户配置部分 ---
    # 1. 您的数据集根目录，例如 'iflytek_phone_ucar'
    #    请确保这个目录与此脚本 (generate_calibration_dataset.py) 在同一级。
    YOUR_DATASET_ROOT = '/home/youngqi/Downloads/ultralytics-main/datasets/iflytek_phone_ucar/'

    # 2. 选择要使用的子文件夹：'train' 或 'val'
    #    通常用于校准的是 'val' (验证集)。
    SUBSET_TO_USE = 'val'

    # 3. 您希望随机选择的图片数量
    #    RKNN 官方建议校准数据集图片数量不低于 300 张，越多越有代表性，但也要适度。
    NUM_IMAGES_FOR_CALIBRATION = 600 

    # 4. 输出的 dataset.txt 文件名
    #    这个文件将包含选中的图片路径列表。
    OUTPUT_DATASET_FILE = 'dataset.txt'
    # --- 配置结束 ---

    create_calibration_dataset_txt(
        base_data_dir=YOUR_DATASET_ROOT,
        subset_folder=SUBSET_TO_USE,
        num_images=NUM_IMAGES_FOR_CALIBRATION,
        output_file=OUTPUT_DATASET_FILE
    )